"""
统一处理器模块

集成记忆、意图识别和处理链，提供完整的处理流程
"""
import asyncio
from datetime import datetime
import logging
from typing import Callable, AsyncGenerator, Dict, Any

from langchain.schema import HumanMessage, AIMessage

from Langchain_utils.chains import create_knowledge_chain, create_general_chain
from Langchain_utils.config import DEFAULT_SESSION_ID
from Langchain_utils.intent import detect_intent
from Langchain_utils.memory import create_memory, format_chat_history
from RAGFlow_utils.query_enhancer import enhance_query


def create_unified_processor(
        system_prompt: str = "你是一个智能助手，可以根据需要查询知识库或直接回答用户问题，并且只能用markdown格式返回数据。",
        session_id: str = DEFAULT_SESSION_ID
) -> Callable[[str], AsyncGenerator[Dict[str, Any], None]]:
    """
    创建统一处理器
    :param system_prompt: 系统提示词
    :param session_id: 会话ID，用于MongoDB存储
    :return:
        Callable: 处理函数，接受问题并返回包含回答和元数据的字典
    """
    # 创建记忆对象，使用指定的会话ID
    memory = create_memory(session_id=session_id)

    # 创建处理链
    knowledge_chain = create_knowledge_chain(system_prompt)
    general_chain = create_general_chain(system_prompt)

    # 处理统计
    stats = {
        "total_queries": 0,
        "knowledge_base_queries": 0,
        "general_chat_queries": 0
    }

    async def processor(question: str) -> AsyncGenerator[Dict[str, Any], None]:
        """
        处理用户问题
        :param question: 用户问题
        :return: 包含回答和元信息的字典
        """
        stats["total_queries"] += 1
        # 加载对话历史
        memory_variables = memory.load_memory_variables({})
        chat_history = memory_variables.get("chat_history", [])

        # 格式化对话历史
        formatted_history = format_chat_history(chat_history)
        # 进行意图分类
        intent = detect_intent(question, formatted_history)
        # 初始化元数据
        metadata = {

            "intent": intent,
            "original_query": question,
            "session_id": session_id

        }

        # 如果是知识库查询，预先计算增强查询（仅用于元数据）
        if intent == "knowledge_base":
            enhanced_query = enhance_query(question, formatted_history)
            metadata["enhanced_query"] = enhanced_query
            if enhanced_query != question:
                metadata["query_enhanced"] = True
            else:
                metadata["query_enhanced"] = False

        # 发送开始元数据
        yield {
            "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "type": "metadata",
            "data": metadata
        }

        # 根据意图选择处理方式并流式输出
        # 根据意图选择处理方式并流式输出
        full_answer = ""
        if intent == "knowledge_base":
            # 知识库处理
            stats["knowledge_base_queries"] += 1
            metadata["processor"] = "knowledge_base"

        else:
            # 一般对话处理
            stats["general_chat_queries"] += 1
            metadata["processor"] = "general_chat"

        async for chunk in knowledge_chain(question, formatted_history):
            full_answer += chunk
            yield {
                "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "type": "content",
                "data": full_answer
            }
            #停顿1s
            await asyncio.sleep(1)

        # 保存完整对话到记忆
        input_message = HumanMessage(content=question)
        output_message = AIMessage(content=full_answer)

        try:
            memory.chat_memory.add_messages([input_message, output_message])
        except Exception as e:
            logging.error(f"将会话历史保存到MongoDB时出错: {e}")

        # 记录查询统计
        logging.info(
            f"处理统计 | 会话ID: {session_id} | 总查询: {stats['total_queries']} | "
            f"知识库查询: {stats['knowledge_base_queries']} | "
            f"一般查询: {stats['general_chat_queries']}"
        )

        # 发送结束标记
        yield {
            "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "type": "end",
            "data": {
                "full_answer": full_answer,
                "metadata": metadata
            }
        }

    # 添加实用函数
    def get_stats() -> Dict[str, int]:
        """获取处理统计信息"""
        return stats.copy()

    def clear_memory() -> None:
        """清除对话记忆"""
        memory.clear()
        logging.info(f"会话 {session_id} 的对话记忆已清除")

    # 将实用函数附加到处理器函数
    processor.get_stats = get_stats
    processor.clear_memory = clear_memory

    return processor
